# !  本文件赋值处理数据相关函数
#. Python标准库,提供基础函数操作 
import os
#. 数字库
import numpy as np
#. 识别XML文件库
import xml.etree.ElementTree as ET
#. 跨平台计算机视觉库
import cv2

#! 常量区
#> 昆虫正确的名字标注等信息
INSECT_NAMES = ['Boerner', 'Leconte', 'Linnaeus', 
                'acuminatus', 'armandi', 'coleoptera', 'linnaeus']


#! 函数区

#> 将昆虫信息转为对应数字格式的JSON字典
#@ 返回值格式: JSON字典格式; 
def get_insect_names():
    insect_category2id = {}
    for i, item in enumerate(INSECT_NAMES):
        insect_category2id[item] = i
    return insect_category2id
#? 返回的数据格式表现为：
"""
实例: {
  'Boerner': 0, 
  'Leconte': 1, 
  'Linnaeus': 2, 
  'acuminatus': 3, 
  'armandi': 4, 
  'coleoptera': 5, 
  'linnaeus': 6
  }
"""


#>  训练数据集的标注数据
#@ cname2cid: 昆虫信息的字典
#@ datadir:  数据文件路径
def get_annotations(cname2cid, datadir):
    filenames = os.listdir(os.path.join(datadir, 'annotations', 'xmls'))
    records = []
    ct = 0
    for fname in filenames:
        fid = fname.split('.')[0]
        fpath = os.path.join(datadir, 'annotations', 'xmls', fname)
        img_file = os.path.join(datadir, 'images', fid + '.jpeg')
        tree = ET.parse(fpath)

        if tree.find('id') is None:
            im_id = np.array([ct])
        else:
            im_id = np.array([int(tree.find('id').text)])

        objs = tree.findall('object')
        im_w = float(tree.find('size').find('width').text)
        im_h = float(tree.find('size').find('height').text)
        gt_bbox = np.zeros((len(objs), 4), dtype=np.float32)
        gt_class = np.zeros((len(objs), ), dtype=np.int32)
        is_crowd = np.zeros((len(objs), ), dtype=np.int32)
        difficult = np.zeros((len(objs), ), dtype=np.int32)
        for i, obj in enumerate(objs):
            cname = obj.find('name').text
            gt_class[i] = cname2cid[cname]
            _difficult = int(obj.find('difficult').text)
            x1 = float(obj.find('bndbox').find('xmin').text)
            y1 = float(obj.find('bndbox').find('ymin').text)
            x2 = float(obj.find('bndbox').find('xmax').text)
            y2 = float(obj.find('bndbox').find('ymax').text)
            x1 = max(0, x1)
            y1 = max(0, y1)
            x2 = min(im_w - 1, x2)
            y2 = min(im_h - 1, y2)
            # 这里使用xywh格式来表示目标物体真实框
            gt_bbox[i] = [(x1+x2)/2.0 , (y1+y2)/2.0, x2-x1+1., y2-y1+1.]
            is_crowd[i] = 0
            difficult[i] = _difficult

        voc_rec = {
            'im_file': img_file,
            'im_id': im_id,
            'h': im_h,
            'w': im_w,
            'is_crowd': is_crowd,
            'gt_class': gt_class,
            'gt_bbox': gt_bbox,
            'gt_poly': [],
            'difficult': difficult
            }
        if len(objs) != 0:
            records.append(voc_rec)
        ct += 1
    return records
#? 返回的数据格式表现为：
"""
实例: {
    数据长度是： 1693
    数据类型是： <class 'list'>
    数据下标 0 的内容是： 
        {
            'im_file': './work/insects/train/images/1985.jpeg', 
            'im_id': array([0]), 
            'h': 1258.0, 
            'w': 1258.0, 
            'is_crowd': array([0, 0, 0, 0, 0, 0], dtype=int32), 
            'gt_class': array([4, 5, 5, 2, 3, 1], dtype=int32), 
            'gt_bbox': array([
                            [837.5, 810.5,  60. ,  74. ],
                            [736.5, 664. ,  52. ,  81. ],
                            [690. , 594. ,  63. ,  37. ],
                            [596. , 685.5,  91. ,  78. ],
                            [917.5, 606.5,  54. ,  82. ],
                            [634.5, 854.5,  88. , 152. ]
                            ], dtype=float32), 
            'gt_poly': [], 
            'difficult': array([0, 0, 0, 0, 0, 0], dtype=int32)}
"""


###! 数据读取

#> 将传进来的真实框的坐标和真实物体类型全部修改为矩阵为50*4的规模
def get_bbox(gt_bbox, gt_class):
    MAX_NUM = 50
    # zeros :返回：返回来一个给定形状和类型的用0填充的数组；用法：zeros(shape, dtype=float, order='C')
    gt_bbox2 = np.zeros((MAX_NUM, 4))
    gt_class2 = np.zeros((MAX_NUM,))
    for i in range(len(gt_bbox)):
        gt_bbox2[i, :] = gt_bbox[i, :]
        gt_class2[i] = gt_class[i]
        if i >= MAX_NUM:
            break
    return gt_bbox2, gt_class2
#？ 返回的数据格式表现为：
"""

[[542.5 652.5 140.  150. ]
 [885.  572.  127.  135. ]
 [648.5 811.5  84.   62. ]
 [798.5 821.   86.   71. ]
 [667.5 521.   88.   67. ]
 [  0.    0.    0.    0. ]
 [  0.    0.    0.    0. ]
 [  0.    0.    0.    0. ]
 其余都是4列为0的数组，
 真实的标记数组是5个，但是真实的图片中需要标注的真实物体框可能远不止5个，所以扩充为50个就怎样都够用了。大概就是这个想法
# 对于一般的检测任务来说，一张图片上往往会有多个目标物体
# 设置参数MAX_NUM = 50， 即一张图片最多取50个真实框；如果真实
# 框的数目少于50个，则将不足部分的gt_bbox, gt_class和gt_score的各项数值全设置为0

"""


#> 返回图片数据的数据，
def get_img_data_from_file(record):
    #. 图片路径
    im_file = record['im_file']
    #. 图片高
    h = record['h']
    #. 图片宽
    w = record['w']
    #. 
    is_crowd = record['is_crowd']
    #. 虫子类别
    gt_class = record['gt_class']
    #. 
    gt_bbox = record['gt_bbox']
    #. 
    difficult = record['difficult']

    img = cv2.imread(im_file)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # 检查记录中的h和w是否等于从img读取的值
    assert img.shape[0] == int(h), \
             "记录（{}）和img文件（{}）中{}的图像高度不一致".format(
               im_file, h, img.shape[0])

    assert img.shape[1] == int(w), \
             "记录（{}）和img文件（{}）中{}的图像宽度不一致".format(
               im_file, w, img.shape[1])

    gt_boxes, gt_labels = get_bbox(gt_bbox, gt_class)

    # gt_bbox 用相对值
    gt_boxes[:, 0] = gt_boxes[:, 0] / float(w)
    gt_boxes[:, 1] = gt_boxes[:, 1] / float(h)
    gt_boxes[:, 2] = gt_boxes[:, 2] / float(w)
    gt_boxes[:, 3] = gt_boxes[:, 3] / float(h)
  
    return img, gt_boxes, gt_labels, (h, w)

#? 返回的数据格式表现为：
"""
    图像数据img 
        数据矩阵式： (1258, 1258, 3)
        长度：      1258
        数据类型：  <class 'numpy.ndarray'>
    真实框坐标gt_boxes
        数据矩阵式： (50, 4)
        长度：      50
        数据类型：  <class 'numpy.ndarray'>
    真实框包含的物体类别gt_labels
        数据矩阵式： (50, )
        长度：      50
        数据类型：  <class 'numpy.ndarray'>
    图像尺寸scales
        长度：      2
        数据类型：  <class 'tuple'>
        数据内容：  (1258.0, 1258.0)

"""
